package org.wenshu.ai.modular.chat.controller;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.annotation.Resource;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.*;
import org.wenshu.ai.config.AiModelConfig;
import org.wenshu.ai.modular.chat.provider.core.AiModelService;
import org.wenshu.ai.modular.datasource.entity.AiDatasourceConfig;
import org.wenshu.ai.modular.datasource.param.AiDatasourceConfigIdParam;
import org.wenshu.ai.modular.datasource.service.AiDatasourceConfigService;
import org.wenshu.ai.modular.traindata.param.TrainDataAddParam;
import org.wenshu.ai.modular.traindata.service.TrainDataService;
import org.wenshu.common.pojo.CommonResult;

import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.DriverManager;
import java.sql.ResultSet;

import static java.time.Duration.ofSeconds;

@Tag(name = "ai数据库控制器")
@RestController
@Validated
public class AiDBController {

    @Resource
    private AiDatasourceConfigService aiDatasourceConfigService;

    @Resource
    private TrainDataService trainDataService;

    @Resource
    private AiModelConfig aiModelConfig;

    @Resource
    private AiModelService aiModelService;

    @Operation(summary = "获取数据库表结构并写入向量库")
    @GetMapping("/ai/db/schema/embed")
    public CommonResult<String> embedDatabaseSchema(@Parameter(description = "数据源id") @RequestParam("datasourceId") String datasourceId) {
        try {
            // 查询数据源配置
            AiDatasourceConfigIdParam idParam = new AiDatasourceConfigIdParam();
            idParam.setId(datasourceId);
            AiDatasourceConfig config = aiDatasourceConfigService.detail(idParam);

            if (!"MYSQL".equalsIgnoreCase(config.getProvider())) {
                return CommonResult.error("目前仅支持MySQL数据源");
            }

            // 获取数据库连接
            Connection conn = DriverManager.getConnection(config.getUrl(), config.getUsername(), config.getPassword());
            DatabaseMetaData metaData = conn.getMetaData();

            // 获取所有表
            ResultSet tables = metaData.getTables(config.getDatabaseName(), null, "%", new String[]{"TABLE"});
            
            while (tables.next()) {
                String tableName = tables.getString("TABLE_NAME");
                String remarks = tables.getString("REMARKS"); // 备注
                StringBuilder schemaText = new StringBuilder();
                schemaText.append("The following columns are in the ")
                        .append(tableName).append(" table ").append(" remarks ").append(remarks)
                        .append(" in the ").append(config.getDatabaseName()).append(" database:\n");
                schemaText.append("| Column Name | Data Type | Column Comment |\n");
                schemaText.append("|------------|-----------|----------------|\n");
                // 获取表的所有列
                ResultSet columns = metaData.getColumns(config.getDatabaseName(), null, tableName, "%");


                while (columns.next()) {
                    String columnName = columns.getString("COLUMN_NAME");
                    String dataType = columns.getString("TYPE_NAME");
                    String columnComment = columns.getString("REMARKS");
                    schemaText.append("| ").append(columnName).append(" | ")
                              .append(dataType).append(" | ")
                              .append(columnComment != null ? columnComment : "").append(" |\n");
                }

                // 保存到TrainData表
                TrainDataAddParam trainDataAddParam = new TrainDataAddParam();
                trainDataAddParam.setDataType("ddl");
                trainDataAddParam.setDdlText(schemaText.toString());
                trainDataService.add(trainDataAddParam);

                columns.close();
            }

            tables.close();
            conn.close();

            return CommonResult.ok("数据库表结构已成功写入向量库");

        } catch (Exception e) {
            return CommonResult.error("处理失败: " + e.getMessage());
        }
    }


}
